#include <stdlib.h>
#include <string.h>
#include "rarray.h"
#include "generator.h"
#include "parser_unref.h"
#include "redintobject.h"
#include "redmapobject.h"
#include "redboolobject.h"
#include "rednullobject.h"
#include "redcodeobject.h"
#include "redarrayobject.h"
#include "redblockobject.h"
#include "redfloatobject.h"
#include "redstringobject.h"

/**************************** generator prototypes ****************************/

/* language structures */

static void generate_if(generator_t *self, if_t *node);
static void generate_for(generator_t *self, for_t *node);
static void generate_try(generator_t *self, try_t *node);
static void generate_while(generator_t *self, while_t *node);
static void generate_switch(generator_t *self, switch_t *node);
static void generate_compond(generator_t *self, compond_t *node);
static void generate_function(generator_t *self, function_t *node);

/* statements */

static void generate_assign(generator_t *self, assign_t *node);
static void generate_delete(generator_t *self, delete_t *node);
static void generate_augment(generator_t *self, augment_t *node);

/* control flows */

static void generate_break(generator_t *self, break_t *node);
static void generate_return(generator_t *self, return_t *node);
static void generate_continue(generator_t *self, continue_t *node);

/* generic statement */

static void generate_statement(generator_t *self, statement_t *node);

/* assign targets */

static void generate_name(generator_t *self, name_t *node);
static void generate_attr(generator_t *self, attr_t *node);
static void generate_index(generator_t *self, index_t *node);
static void generate_invoke(generator_t *self, invoke_t *node);
static void generate_rvalue(generator_t *self, rvalue_t *node);

/* expressions */

static void generate_expr(generator_t *self, expr_t *node);
static void generate_unit(generator_t *self, unit_t *node);
static void generate_const(generator_t *self, const_t *node);
static void generate_range(generator_t *self, range_t *node);

/**************************** generator implementations ****************************/

static inline opcode_t augment_map(const char *op)
{
	if (!strcmp(op, "+="))		return AUG_ADD;
	if (!strcmp(op, "-="))		return AUG_SUB;
	if (!strcmp(op, "*="))		return AUG_MUL;
	if (!strcmp(op, "/="))		return AUG_DIV;
	if (!strcmp(op, "%="))		return AUG_MOD;
	if (!strcmp(op, "**="))		return AUG_POW;

	if (!strcmp(op, "|="))		return AUG_BIT_OR;
	if (!strcmp(op, "&="))		return AUG_BIT_AND;
	if (!strcmp(op, "^="))		return AUG_BIT_XOR;
	if (!strcmp(op, "<<="))		return AUG_LSHIFT;
	if (!strcmp(op, ">>="))		return AUG_RSHIFT;

	/* this should never happens */
	abort();
}

static inline opcode_t operator_map(const char *op)
{
	if (!strcmp(op, "<"))		return LE;
	if (!strcmp(op, ">"))		return GE;
	if (!strcmp(op, "<="))		return LEQ;
	if (!strcmp(op, ">="))		return GEQ;
	if (!strcmp(op, "=="))		return EQ;
	if (!strcmp(op, "!="))		return NEQ;

	if (!strcmp(op, "+"))		return ADD;
	if (!strcmp(op, "-"))		return SUB;
	if (!strcmp(op, "*"))		return MUL;
	if (!strcmp(op, "/"))		return DIV;
	if (!strcmp(op, "%"))		return MOD;
	if (!strcmp(op, "**"))		return POW;

	if (!strcmp(op, "|"))		return BIT_OR;
	if (!strcmp(op, "&"))		return BIT_AND;
	if (!strcmp(op, "^"))		return BIT_XOR;
	if (!strcmp(op, "<<"))		return LSHIFT;
	if (!strcmp(op, ">>"))		return RSHIFT;

	if (!strcmp(op, "or"))		return BOOL_OR;
	if (!strcmp(op, "and"))		return BOOL_AND;
	if (!strcmp(op, "xor"))		return BOOL_XOR;

	/* this should never happens */
	abort();
}

static inline char find_name(generator_t *self, rstring_t *name)
{
	for (int i = 0; i < rarray_count(self->names); i++)
		if (rstring_equals_rs(name, rarray_item_at(self->names, i, rstring_t)))
			return 1;

	return 0;
}

/* language structures */

#define PC				(self->current_buffer->pc)
#define emit(op)		bytecode_buffer_emit(self->current_buffer, op, node->node.lineno, node->node.position)

static void generate_if(generator_t *self, if_t *node)
{
	generate_expr(self, node->expr);
	bytecode_t *branch = emit(BRFALSE);
	generate_statement(self, node->pos);

	if (node->neg == NULL)
	{
		branch->at_offset = PC;
		return;
	}

	bytecode_t *skip = emit(BR);
	branch->at_offset = PC;
	generate_statement(self, node->neg);
	skip->at_offset = PC;
}

static void _patch_branches(rarray_t *self, void *data, void *context)
{
	(void)self;
	((bytecode_t *)data)->at_offset = *(int *)context;
}

static void generate_for(generator_t *self, for_t *node)
{
	generate_statement(self, node->init);
	int start = PC;
	generate_expr(self, node->condition);
	bytecode_t *branch = emit(BRFALSE);

	rarray_context_t *breaks;
	rarray_context_t *continues;

	rarray_save_context(&(self->breaks), &breaks);
	rarray_save_context(&(self->continues), &continues);

	generate_statement(self, node->body);
	rarray_foreach(&(self->continues), _patch_branches, &PC);
	rarray_restore_context(&(self->continues), &continues);

	generate_statement(self, node->step);
	emit(BR)->at_offset = start;
	branch->at_offset = PC;

	rarray_foreach(&(self->breaks), _patch_branches, &PC);
	rarray_restore_context(&(self->breaks), &breaks);
}

static void _patch_excepts(rarray_t *self, void *data, void *context)
{
	((bytecode_t *)data)->at_offset = *(int *)context;
}

typedef struct _ge_ctx_t
{
	RedObject		*block;
	rarray_t		 excepts;
	generator_t		*generator;
} ge_ctx_t;

static void _generate_except(rarray_t *self, void *data, void *context)
{
	except_t *exc = (except_t *)data;
	ge_ctx_t *ctx = (ge_ctx_t *)context;
	RedBlockEntry *entry = RedBlock_Append(ctx->block);

	for (int i = 0; i < rarray_count(exc->types); i++)
		rarray_append(&(entry->e_types), RedString_FromRString(rarray_item_at(exc->types, i, rstring_t)));

	entry->e_pc = ctx->generator->current_buffer->pc;
	entry->e_name = exc->name ? RedString_FromRString(exc->name) : NULL;

	generate_statement(ctx->generator, exc->body);
	rarray_append(&(ctx->excepts), bytecode_buffer_emit(ctx->generator->current_buffer, BR, exc->node.lineno, exc->node.position));
}

static void generate_try(generator_t *self, try_t *node)
{
	RedObject *block = RedBlock_Create();

	emit(PUSH_BLOCK)->at_object = block;
	generate_statement(self, node->body);

	ge_ctx_t context;
	bytecode_t *branch = emit(BR);

	context.block = block;
	context.generator = self;
	rarray_create(&(context.excepts), NULL);
	rarray_foreach(&(node->excepts), _generate_except, &context);

	if (node->generic)
	{
		RedBlock_Generic(block) = PC;
		generate_statement(self, node->generic);
	}

	branch->at_offset = PC;
	rarray_foreach(&(context.excepts), _patch_excepts, &PC);

	if (node->finally)
	{
		RedBlock_Finally(block) = PC;
		generate_statement(self, node->finally);
	}

	emit(POP_BLOCK);
}

static void generate_while(generator_t *self, while_t *node)
{
	int start = PC;
	generate_expr(self, node->expr);
	bytecode_t *branch = emit(BRFALSE);

	rarray_context_t *breaks;
	rarray_context_t *continues;

	rarray_save_context(&(self->breaks), &breaks);
	rarray_save_context(&(self->continues), &continues);

	generate_statement(self, node->body);
	rarray_foreach(&(self->continues), _patch_branches, &start);
	rarray_restore_context(&(self->continues), &continues);

	emit(BR)->at_offset = start;
	branch->at_offset = PC;

	rarray_foreach(&(self->breaks), _patch_branches, &PC);
	rarray_restore_context(&(self->breaks), &breaks);
}

static void generate_switch(generator_t *self, switch_t *node)
{
	rarray_t patch;
	rarray_t pending;

	generate_expr(self, node->expr);
	rarray_create(&patch, NULL);
	rarray_create(&pending, NULL);

	for (int i = 0; i < rarray_count(node->cases); i++)
	{
		case_t *item = rarray_item_at(node->cases, i, case_t);

		for (int j = 0; j < rarray_count(item->ranges); j++)
		{
			emit(DUP);
			generate_range(self, rarray_item_at(item->ranges, j, range_t));

			if (j != rarray_count(item->ranges) - 1)
				rarray_append(&pending, emit(BRTRUE));
		}

		bytecode_t *branch = emit(BRFALSE);

		rarray_foreach(&pending, _patch_branches, &PC);
		rarray_clear(&pending);
		generate_statement(self, item->body);
		rarray_append(&patch, emit(BR));
		branch->at_offset = PC;
	}

	if (node->other)
		generate_statement(self, node->other);

	rarray_foreach(&patch, _patch_branches, &PC);
	emit(DROP);

	rarray_unref(&patch);
	rarray_unref(&pending);
}

static void _generate_statement(rarray_t *self, void *data, void *context)
{
	(void)self;
	generate_statement((generator_t *)context, (statement_t *)data);
}

static void generate_compond(generator_t *self, compond_t *node)
{
	rarray_foreach(&(node->statements), _generate_statement, self);
}

typedef struct _la_ctx_t
{
	function_t			*node;
	rarray_t			*names;
	bytecode_buffer_t	*buffer;
} la_ctx_t;

static void generate_function(generator_t *self, function_t *node)
{
	bytecode_buffer_t *old = self->current_buffer;
	bytecode_buffer_t *new = bytecode_buffer_create();

	rarray_context_t *free;
	rarray_context_t *names;
	self->current_buffer = new;
	rarray_save_context(&(self->free), &free);
	rarray_save_context(&(self->names), &names);

	/* register argument as locals */
	for (int i = 0; i < rarray_count(node->args); i++)
		rarray_append(&(self->names), rarray_item_at(node->args, i, rstring_t));

	/* generate bytecodes for function body */
	generate_statement(self, node->body);

	Red_REF(Red_Null);
	emit(LOAD_CONST)->at_object = Red_Null;
	emit(POP_RETURN);

	/* link together, for optimization of cache access */
	RedObject *code = bytecode_buffer_assemble(new, self->fname, self->source, self->module);

	/* register argument names */
	RedCode_Args(code) = RedArray_CreatePresized(rarray_count(node->args));
	RedCode_Frees(code) = RedArray_CreatePresized(rarray_count(self->free));

	/* if variable argument is present */
	if (node->vararg)
		RedCode_VarArg(code) = RedString_FromRString(node->vararg);

	/* transfer arguments into code object */
	for (int i = 0; i < rarray_count(node->args); i++)
		RedArray_Items(RedCode_Args(code))[i] = RedString_FromRString(rarray_item_at(node->args, i, rstring_t));

	/* transfer free variables into code object */
	for (int i = 0; i < rarray_count(self->free); i++)
		RedArray_Items(RedCode_Frees(code))[i] = RedString_FromRString(rarray_item_at(self->free, i, rstring_t));

	self->current_buffer = old;
	rarray_restore_context(&(self->free), &free);
	rarray_restore_context(&(self->names), &names);
	bytecode_buffer_release(&new);

	emit(LOAD_CONST)->at_object = code;
	emit(MAKE_FUNCTION)->at_offset = rarray_count(node->args);

	/* anonymous function support */
	if (node->name)
	{
		rarray_append(&(self->names), node->name);
		emit(STOR_OBJECT)->at_object = RedString_FromRString(node->name);
	}
}

/* statements */

static void generate_assign(generator_t *self, assign_t *node)
{
	switch (node->type)
	{
		case tv_attr:
		{
			generate_rvalue(self, node->tv_attr->value);
			generate_expr(self, node->expr);
			emit(SET_ATTR)->at_object = RedString_FromRString(node->tv_attr->attribute);
			break;
		}

		case tv_name:
		{
			generate_expr(self, node->expr);
			rarray_append(&(self->names), node->tv_name->name);
			emit(STOR_OBJECT)->at_object = RedString_FromRString(node->tv_name->name);
			break;
		}

		case tv_index:
		{
			generate_rvalue(self, node->tv_index->value);
			generate_expr(self, node->tv_index->expr);
			generate_expr(self, node->expr);
			emit(SET_ITEM);
			break;
		}
	}
}

static void generate_delete(generator_t *self, delete_t *node)
{
	switch (node->type)
	{
		case tv_attr:
		{
			generate_rvalue(self, node->tv_attr->value);
			emit(DEL_ATTR)->at_object = RedString_FromRString(node->tv_attr->attribute);
			break;
		}

		case tv_name:
		{
			emit(DEL_OBJECT)->at_object = RedString_FromRString(node->tv_name->name);
			break;
		}

		case tv_index:
		{
			generate_rvalue(self, node->tv_index->value);
			generate_expr(self, node->tv_index->expr);
			emit(DEL_ITEM);
			break;
		}
	}
}

static void generate_augment(generator_t *self, augment_t *node)
{
	switch (node->type)
	{
		case tv_attr:
		{
			generate_rvalue(self, node->tv_attr->value);
			emit(DUP);
			emit(GET_ATTR)->at_object = RedString_FromRString(node->tv_attr->attribute);
			break;
		}

		case tv_name:
		{
			if (!find_name(self, node->tv_name->name))
				rarray_append(&(self->free), node->tv_name->name);

			emit(LOAD_OBJECT)->at_object = RedString_FromRString(node->tv_name->name);
			break;
		}

		case tv_index:
		{
			generate_rvalue(self, node->tv_index->value);
			generate_expr(self, node->tv_index->expr);
			emit(DUP2);
			emit(GET_ITEM);
			break;
		}
	}

	generate_expr(self, node->expr);
	emit(augment_map(node->operator->content));

	switch (node->type)
	{
		case tv_attr:
		{
			emit(SET_ATTR)->at_object = RedString_FromRString(node->tv_attr->attribute);
			break;
		}

		case tv_name:
		{
			rarray_append(&(self->names), node->tv_name->name);
			emit(STOR_OBJECT)->at_object = RedString_FromRString(node->tv_name->name);
			break;
		}

		case tv_index:
		{
			emit(SET_ITEM);
			break;
		}
	}
}

/* control flows */

static void generate_break(generator_t *self, break_t *node)
{
	(void)node;
	rarray_append(&(self->breaks), emit(BR));
}

static void generate_return(generator_t *self, return_t *node)
{
	generate_expr(self, node->expr);
	emit(POP_RETURN);
}

static void generate_continue(generator_t *self, continue_t *node)
{
	(void)node;
	rarray_append(&(self->continues), emit(BR));
}

/* generic statement */

static void generate_statement(generator_t *self, statement_t *node)
{
#define dispatch(t)	case st_##t: generate_##t(self, node->st_##t); break;

	switch (node->type)
	{
		dispatch(if)
		dispatch(for)
		dispatch(try)
		dispatch(while)
		dispatch(switch)
		dispatch(compond)
		dispatch(function)

		dispatch(assign)
		dispatch(delete)
		dispatch(invoke)
		dispatch(augment)

		dispatch(break)
		dispatch(return)
		dispatch(continue)
	}

#undef dispatch
}

/* assign targets */

static void generate_name(generator_t *self, name_t *node)
{
	if (!find_name(self, node->name))
		rarray_append(&(self->free), node->name);

	RedObject *name = RedString_FromRString(node->name);
	RedObject *object = RedMap_GetString(self->consts, name);

	if (!object)
	{
		/* not a built-in object, load it by name */
		emit(LOAD_OBJECT)->at_object = name;
	}
	else
	{
		/* built-in object can be loaded directly as const */
		Red_REF(object);
		Red_UNREF(name);
		emit(LOAD_CONST)->at_object = object;
	}
}

static void generate_attr(generator_t *self, attr_t *node)
{
	generate_rvalue(self, node->value);
	emit(GET_ATTR)->at_object = RedString_FromRString(node->attribute);
}

static void generate_index(generator_t *self, index_t *node)
{
	generate_rvalue(self, node->value);
	generate_expr(self, node->expr);
	emit(GET_ITEM);
}

static void generate_invoke(generator_t *self, invoke_t *node)
{
	if (!node->variable)
	{
		for (int i = rarray_count(node->args) - 1; i >= 0; i--)
			generate_expr(self, rarray_item_at(node->args, i, expr_t));

		generate_rvalue(self, node->value);
		emit(INVOKE)->at_offset = rarray_count(node->args);
	}
	else
	{
		emit(START_ARGS);

		for (int i = rarray_count(node->args) - 1; i >= 0; i--)
		{
			generate_expr(self, rarray_item_at(node->args, i, expr_t));
			if (rarray_item_at(node->expand, i, intptr_t)) emit(EXPAND);
		}

		generate_rvalue(self, node->value);
		emit(INVOKE_VARG);
	}

	if (node->discard)
		emit(DROP);
}

static void generate_rvalue(generator_t *self, rvalue_t *node)
{
	switch (node->type)
	{
		case rv_attr:	generate_attr(self, node->rv_attr);		break;
		case rv_name:	generate_name(self, node->rv_name);		break;
		case rv_unit:	generate_unit(self, node->rv_unit);		break;
		case rv_const:	generate_const(self, node->rv_const);	break;
		case rv_index:	generate_index(self, node->rv_index);	break;
		case rv_invoke:	generate_invoke(self, node->rv_invoke);	break;
	}
}

/* expressions */

static void generate_expr(generator_t *self, expr_t *node)
{
	switch (node->el_type)
	{
		case e_expr:	generate_expr(self, node->el_expr);		break;
		case e_rvalue:	generate_rvalue(self, node->el_rvalue);	break;
	}

	if (node->operator)
	{
		opcode_t opcode = operator_map(node->operator->content);
		bytecode_t *branch = NULL;

		switch (opcode)
		{
			case BOOL_OR:
			{
				emit(DUP);
				branch = emit(BRTRUE);
				break;
			}

			case BOOL_AND:
			{
				emit(DUP);
				branch = emit(BRFALSE);
				break;
			}

			default:
				break;
		}

		switch (node->er_type)
		{
			case e_expr:	generate_expr(self, node->er_expr);		break;
			case e_rvalue:	generate_rvalue(self, node->er_rvalue);	break;
		}

		emit(opcode);

		if (branch)
			branch->at_offset = PC;
	}
}

static void generate_unit(generator_t *self, unit_t *node)
{
	switch (node->type)
	{
		case u_expr:		generate_expr(self, node->u_expr); break;
		case u_unit:		generate_unit(self, node->u_unit); break;
		case u_rvalue:		generate_rvalue(self, node->u_rvalue); break;
		case u_function:	generate_function(self, node->u_function); break;
	}

	if (node->operator)
	{
		if (rstring_equals(node->operator, "+"))			emit(POS);
		else if (rstring_equals(node->operator, "-"))		emit(NEG);
		else if (rstring_equals(node->operator, "~"))		emit(BIT_NOT);
		else if (rstring_equals(node->operator, "not"))		emit(BOOL_NOT);
	}
}

static void generate_const(generator_t *self, const_t *node)
{
	switch (node->type)
	{
		case c_int:		emit(LOAD_CONST)->at_object = RedInt_FromInt(node->c_int);				break;
		case c_float:	emit(LOAD_CONST)->at_object = RedFloat_FromFloat(node->c_float);		break;
		case c_string:	emit(LOAD_CONST)->at_object = RedString_FromRString(node->c_string);	break;
	}
}

static void generate_range(generator_t *self, range_t *node)
{
	generate_expr(self, node->begin);

	if (node->end == NULL)
		emit(DUP);
	else
		generate_expr(self, node->end);

	emit(RANGE);
}

/**************************** generator dispatcher ****************************/

RedObject *generator_generate(compond_t *ast, rstring_t *fname, rstring_t *source, RedObject *module, RedObject *consts)
{
	generator_t generator;
	rarray_create(&(generator.free), NULL);
	rarray_create(&(generator.names), NULL);
	rarray_create(&(generator.breaks), NULL);
	rarray_create(&(generator.continues), NULL);

	generator.fname = fname;
	generator.consts = consts;
	generator.source = source;
	generator.module = module;
	generator.current_buffer = bytecode_buffer_create();

	Red_REF(Red_Null);
	generate_compond(&generator, ast);
	bytecode_buffer_emit(generator.current_buffer, LOAD_CONST, -1, -1)->at_object = Red_Null;
	bytecode_buffer_emit(generator.current_buffer, POP_RETURN, -1, -1);

	RedObject *result = bytecode_buffer_assemble(generator.current_buffer, fname, source, module);

	unref_compond(ast);
	rarray_unref(&(generator.free));
	rarray_unref(&(generator.names));
	rarray_unref(&(generator.breaks));
	rarray_unref(&(generator.continues));
	bytecode_buffer_release(&(generator.current_buffer));
	return result;
}
